Skip to content

Fix multimodal model registrations in nnx_decoders#4320

Merged
copybara-service[bot] merged 1 commit into
mainfrom
hengtaoguo-nnx-fix
Jul 1, 2026
Merged

Fix multimodal model registrations in nnx_decoders#4320
copybara-service[bot] merged 1 commit into
mainfrom
hengtaoguo-nnx-fix

Conversation

@hengtaoguo

@hengtaoguo hengtaoguo commented Jul 1, 2026

Copy link
Copy Markdown
Collaborator

Description

With the migration to nnx_decoders.py, some recent changes in multimodal were not covered. This PR fixes the issue by adding some new model names registration (along with code-style format updates).

#4263 #4293

Tests

# Checkpoint conversion
python -m maxtext.checkpoint_conversion.to_maxtext src/maxtext/configs/base.yml model_name=qwen3-vl-2b base_output_directory=gs://hengtaoguo-maxtext-logs/checkpoints/qwen3-vl-2b/unscanned/2026-07-01-11-55 scan_layers=false hf_access_token=xxx weight_dtype=bfloat16 hardware=cpu skip_jax_distributed_system=True checkpoint_storage_use_ocdbt=False checkpoint_storage_use_zarr3=False --eager_load_method=safetensors --lazy_load_tensors=False

# Decode
python -m maxtext.inference.decode maxtext/configs/base.yml tokenizer_type=huggingface model_name=qwen3-vl-2b tokenizer_path=Qwen/Qwen3-VL-2B-Instruct load_parameters_path=gs://hengtaoguo-maxtext-logs/checkpoints/qwen3-vl-2b/unscanned/2026-07-01-11-55/0/items per_device_batch_size=1 run_name=ht_test scan_layers=false use_multimodal=true prompt=\'Describe\ this\ image\' image_path=\'tests/assets/test_image.jpg\' max_prefill_predict_length=512 max_target_length=576 ici_tensor_parallelism=4 override_model_config=true attention=\'dot_product\' hf_access_token=xxx

# Result
Input `<|im_start|>user
<|vision_start|><|image_pad|><|vision_end|>Describe this image<|im_end|>
<|im_start|>assistant
` -> `This is a panoramic view of the Seattle skyline on a bright, sunny day. The image is taken from a high vantage point, lookingdown on the city.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@codecov

codecov Bot commented Jul 1, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 0% with 1 line in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/layers/nnx_decoders.py 0.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@hengtaoguo hengtaoguo closed this Jul 1, 2026
@hengtaoguo hengtaoguo force-pushed the hengtaoguo-nnx-fix branch from 22615b5 to a877e45 Compare July 1, 2026 21:27
Comment thread src/maxtext/layers/nnx_decoders.py
@abhinavclemson abhinavclemson self-requested a review July 1, 2026 21:48
@copybara-service copybara-service Bot merged commit 1300488 into main Jul 1, 2026
99 of 105 checks passed
@copybara-service copybara-service Bot deleted the hengtaoguo-nnx-fix branch July 1, 2026 23:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants